import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.stattools import acf, pacf
import seaborn as sns

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

def analyze_acf_pacf(series, lags=30, title="时间序列ACF/PACF分析"):
    """
    分析时间序列的自相关(ACF)和偏自相关(PACF)
    
    参数:
    series: 时间序列数据，可以是pandas Series或numpy数组
    lags: 要计算的滞后阶数，默认为30
    title: 图表标题
    """
    # 确保输入是pandas Series
    if not isinstance(series, pd.Series):
        series = pd.Series(series)
    
    # 计算ACF和PACF值
    acf_values = acf(series, nlags=lags)
    pacf_values = pacf(series, nlags=lags)
    
    # 创建图形
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))
    fig.suptitle(title, fontsize=16)
    
    # 绘制ACF图
    plot_acf(series, lags=lags, ax=ax1, title="自相关函数(ACF)")
    ax1.set_xlabel("滞后阶数")
    ax1.set_ylabel("自相关系数")
    
    # 绘制PACF图
    plot_pacf(series, lags=lags, ax=ax2, title="偏自相关函数(PACF)")
    ax2.set_xlabel("滞后阶数")
    ax2.set_ylabel("偏自相关系数")
    
    plt.tight_layout(rect=[0, 0, 1, 0.96])  # 调整布局，避免标题重叠
    plt.show()
    
    # 打印ACF和PACF的关键值
    print("ACF值（前10个滞后阶数）:")
    print(pd.DataFrame({"滞后阶数": range(11), "ACF值": acf_values[:11]}).to_string(index=False))
    
    print("\nPACF值（前10个滞后阶数）:")
    print(pd.DataFrame({"滞后阶数": range(11), "PACF值": pacf_values[:11]}).to_string(index=False))
    
    return acf_values, pacf_values

def generate_sample_series(data_type="stationary", n=1000):
    """生成不同类型的示例时间序列数据用于测试"""
    np.random.seed(42)  # 设置随机种子，保证结果可复现
    
    if data_type == "stationary":
        # 平稳时间序列（白噪声）
        return pd.Series(np.random.normal(0, 1, n), name="平稳序列（白噪声）")
    
    elif data_type == "trend":
        # 带有趋势的时间序列
        trend = np.linspace(0, 10, n)
        noise = np.random.normal(0, 1, n)
        return pd.Series(trend + noise, name="带趋势的序列")
    
    elif data_type == "seasonal":
        # 带有季节性的时间序列
        t = np.arange(n)
        seasonal = np.sin(2 * np.pi * t / 24)  # 周期为24的季节性
        noise = np.random.normal(0, 0.3, n)
        return pd.Series(seasonal + noise, name="带季节性的序列")
    
    elif data_type == "ar1":
        # AR(1)过程：x_t = 0.7*x_{t-1} + ε_t
        x = np.zeros(n)
        for t in range(1, n):
            x[t] = 0.7 * x[t-1] + np.random.normal(0, 1)
        return pd.Series(x, name="AR(1)序列")
    
    elif data_type == "ma1":
        # MA(1)过程：x_t = ε_t + 0.6*ε_{t-1}
        eps = np.random.normal(0, 1, n)
        x = np.zeros(n)
        for t in range(1, n):
            x[t] = eps[t] + 0.6 * eps[t-1]
        return pd.Series(x, name="MA(1)序列")

# 示例用法
if __name__ == "__main__":
    # 生成不同类型的示例序列
    series_types = ["stationary", "trend", "seasonal", "ar1", "ma1"]
    gdp = pd.read_excel('data/GDP_不变价_当季同比.xlsx', index_col=0)
    gdp = gdp.reset_index().iloc[:, -1]
    # 选择一个序列类型进行分析（可修改索引）
    selected_type = series_types[3]  # 这里选择AR(1)序列进行分析
    # series = generate_sample_series(selected_type)
    series = gdp
    series.name = 'GDP'
    
    # 绘制序列图
    plt.figure(figsize=(12, 4))
    sns.lineplot(data=series)
    plt.title(f"{series.name}时序图")
    plt.xlabel("时间")
    plt.ylabel("值")
    plt.show()
    
    # 分析ACF和PACF
    acf_vals, pacf_vals = analyze_acf_pacf(series, lags=30, title=f"{series.name}的ACF和PACF分析")
    